#!/usr/bin/env python2
# coding=utf-8
#导入Pyqt必须的库
from PyQt4 import QtGui,QtCore
import time
import sys,os
import sqlite3
import threading
import thread
import re
import json
import hashlib
sys.path.append(os.getcwd()+'/UIlib')
sys.path.append(os.getcwd()+'/SqlInject')
from SqlInject.sqlmap.startapi import apiStart
from SqlInject.BatchTest import startSqlthread
from SqlInject.AutoStart import AutoStart
from data.DataBaseHelper import dataBaseHelper
from UIlib.UIthread import startThread
from UIlib.mainUI import Ui_MainWindow 
from UIlib.URLdata import UrlData
from UIlib.uniqueUrl import uniqueUrl
from UIlib.setting import Ui_Setting_Dialog, _translate
from XSSScaner.scaner import Scaner
class StartQT4(QtGui.QMainWindow):
    logemit = QtCore.pyqtSignal(str)

    def __init__(self, parent=None):
        QtGui.QWidget.__init__(self, parent)
        self.ui = Ui_MainWindow()
        self.ui.setupUi(self)
        self.my_label = QtGui.QLabel(self)
        self.my_label.setText(u'这里显示状态信息')
        self.ui.statusbar.addWidget(self.my_label)
        self.count = 0
        self.status = 0
        self.is_scanning = False
        
        #表格单元格无法修改
        self.ui.result_tableWidget.setEditTriggers(QtGui.QAbstractItemView.NoEditTriggers)
        self.ui.info_Edit.setReadOnly(True)
        self.ui.log_Edit.setReadOnly(True)

        self.ui.stopButton.setDisabled(True)
        
        self.table_set = tableThread()
        self.spider = None
        self.sql_scan = sqlScanThread()
        self.sql_map = sqlmapThread()
        self.xss_scan = Scaner()
        self.sql_result = insertSqlRuseltThread()
        self.ui.result_tableWidget.scrollToBottom()
        self.db_select = dbClass()
        
        self.setting = QtGui.QDialog()
        self.setpyui = Ui_Setting_Dialog()
        self.setpyui.setupUi(self.setting)
        QtCore.QObject.connect(self.setpyui.pushButton_add, QtCore.SIGNAL('clicked()'), self.addXSSPattern)
        QtCore.QObject.connect(self.setpyui.pushButton_delete, QtCore.SIGNAL('clicked()'), self.deleteXSSPattern)
        QtCore.QObject.connect(self.setpyui.pushButton_revise, QtCore.SIGNAL('clicked()'), self.modifyXSSPattern)
        QtCore.QObject.connect(self.setpyui.listWidget_rule, QtCore.SIGNAL('itemClicked(QListWidgetItem *)'), self._itemClicked)
        self.initSetting()
        
        QtCore.QObject.connect(self.ui.startButton,QtCore.SIGNAL("clicked()"), self.startSolt)#开始
        QtCore.QObject.connect(self.ui.cleanButton,QtCore.SIGNAL("clicked()"), self.clearDB)#清除数据
        QtCore.QObject.connect(self.ui.stopButton,QtCore.SIGNAL("clicked()"), self.stopSolt)#停止
        QtCore.QObject.connect(self.ui.exportButton,QtCore.SIGNAL("clicked()"), self.showResult)#显示结果
        QtCore.QObject.connect(self.ui.actionSkin,QtCore.SIGNAL("triggered()"), self.settingSolt)#高级设置
        self.logemit.connect(self.logSet)
        self.sql_scan.logemit.connect(self.logSet)
        self.sql_scan.labemit.connect(self.inforLable)
        self.sql_scan.endemit.connect(self.insertResult)
        self.xss_scan.logemit.connect(self.logSet)
        self.xss_scan.endimit.connect(self.showResult)
        self.sql_result.logemit.connect(self.logSet)
        self.sql_result.endinsert.connect(self.endSql)
        self.sql_map.logemit.connect(self.logSet)
        self.table_set.trigger.connect(self.setTable)
        self.table_set.table.connect(self.setTableRow)
        self.ui.result_tableWidget.itemClicked.connect(self.inforSolt)
        #QtCore.QObject.connect(self,QtCore.SIGNAL("logemit()"), self.logSet)
        
        
        self.table_set.start()
        
    #点击开始事件
    def startSolt(self):
        self.is_scanning = True
        self.switch()
        self.getConfig()
        if UrlData.spiderflage:
            self.startSpider()
        elif UrlData.removflage:
            pass
        elif UrlData.sqlflage:
            self.sql_map.start()
            self.sql_scan.start()

        elif UrlData.xssflage:
            self.xss_scan.start()

     #点击清除所有记录
    def clearDB(self):
        try:
            self.db_select.connectDb()
            self.db_select.delectDb()
            self.db_select.disConnect()
            count_tmp = self.count+1
            self.count = 0
            self.db_select.connectDb()
            lensum = self.ui.result_tableWidget.rowCount()
            for i in range(count_tmp):
                self.ui.result_tableWidget.removeRow(lensum)
                lensum -= 1
            self.logSet(u'数据已清除')
        except:
            self.logSet(u'清除异常')
        
    #点击停止事件
    def stopSolt(self):
        self.is_scanning = False
        self.switch()

        UrlData.end_flag = True
        self.table_set.exit()
        self.logSet(u'正在停止....')
        if self.xss_scan.isRunning():
            self.xss_scan.terminate()
            self.logemit.emit(u'XSS检查已经停止')
        if self.spider and self.spider.is_alive:
            self.spider.stop()
        
    #点击高级设置
    def settingSolt(self):
        self.setting.show()    
        self.logSet(u'高级设置')

        
    def getConfig(self):
        UrlData.spiderflage = self.ui.checkBox_URL.isChecked()
        UrlData.removflage = self.ui.checkBox.isChecked()
        UrlData.sqlflage = self.ui.checkBox_SQL.isChecked()
        UrlData.xssflage = self.ui.checkBox_XSS.isChecked()
        UrlData.depth = self.ui.depth_spinBox.value()
        UrlData.threads = self.ui.thread_spinBox.value()
        UrlData.spidemode = self.ui.comboBox.currentIndex()
        
            
    def startSpider(self):
        """
        爬虫初始化
        """        
  
        if self.ui.URL_Edit.text():
            url = str(self.ui.URL_Edit.text())
        else:     
            self.logemit.emit(u'请输入检测目标')  
            return
        depth = self.ui.depth_spinBox.value()
        thread = self.ui.thread_spinBox.value()
        self.logemit.emit(u'测试开始')
        #开始爬虫
        depth = self.ui.depth_spinBox.value()
        threads = self.ui.thread_spinBox.value()
        if self.spider:
            del self.spider
        self.spider = startThread(url,depth,threads,UrlData.spidemode)
        if UrlData.spidemode == 0:
            text = u'随机'
        elif UrlData.spidemode == 1:
            text = u'深度优先'
        elif UrlData.spidemode == 2:
            text = u'广度优先'
        self.logemit.emit(u'爬虫开始 目标:'+url+" 深度:"+str(depth)+" 线程数:"+str(threads)+"爬虫模式: "+ text)
        self.spider.start()

    def inforSolt(self,item = None):
        """
            显示详细信息
        """      
        se_flag = False
        if item==None:
            return        
        introw = self.ui.result_tableWidget.row(item)
        url_id = self.ui.result_tableWidget.item(introw,0).text()
        self.ui.info_Edit.clear()
        
        self.db_select.connectDb()
        cursor = self.db_select.selectAllInfo(url_id)
        try:
            row = cursor.fetchone()
            sql_re = row[3]
            if sql_re == '1' : 
                sql_re = u'存在安全漏洞'
                se_flag = True
            else:
                sql_re = u'安全'
            try:
                xss_re = row[7]
            except Exception as e:
                xss_re = u'暂无'
            self.ui.info_Edit.insertPlainText(u'目标:\t'+row[1]+u' \n日期:\t' + row[2] + u'\nMethod:\t' + row[4] +u'\nSQL检查结果:\t'+ sql_re +u'\nXSS检查结果:\t'+xss_re)
            if se_flag:
                try:
                    f = open('./data/temp/'+str(hashlib.md5(row[1]).hexdigest())+'.json')
                    djson = json.load(f)
                    dbms = djson[0]['value'][0]['dbms'][0]
                    server_os = djson[0]['value'][0]['os']
                    payload = djson[0]['value'][0]['data']['1']['payload']
                    self.ui.info_Edit.insertPlainText(u'\n操作系统:\t'+unicode(server_os))
                    self.ui.info_Edit.insertPlainText(u'\n数据库:\t'+unicode(dbms))
                    self.ui.info_Edit.insertPlainText(u'\npayload:\t'+unicode(payload))
                except:
                    self.ui.info_Edit.insertPlainText(u'\n暂无详细信息')       
            if xss_re == 'danger':
               try:
                    xss_payload = self.db_select.selectXSSpayload(int(url_id))
                    xss_payload = xss_payload.fetchone()
                    for row in xss_payload:
                        self.ui.info_Edit.insertPlainText(u'\npayload:\t'+unicode(row))
               except Exception as e:
                    self.ui.info_Edit.insertPlainText(u'\nNo information\n' + unicode(e))
               
        except Exception as e:
            
            self.ui.info_Edit.insertPlainText(u'\n目标无扫描结果\n' + unicode(e))
        
        self.db_select.disConnect()
            
    def inforLable(self,item = None):
        """
        底栏状态信息显示
        """
        
        if item == None:
            return
        self.my_label.setText(u'    '+item)
          
    #设置Table中值
    def setTable(self,row,flag = 0):
        self.ui.result_tableWidget.insertRow(self.count)
        self.ui.result_tableWidget.setRowHeight(self.count,20)
        for i in range (7):
            re_flag = False
            if i==0:
                text = row[0]
            elif i==1:
                text = row[2]
            elif i==2:
                text = row[4]
            elif i==3 :
                text = row[3]
                if text == '0' and  flag == 1: text = '无安全漏洞'
                elif not text: text = '未知'
                else:
                    re_flag = True
                    text = '存在风险'
            elif i==4:
                if len(row) >= 7:
                    text = row[7]
                    if text == 'danger' : 
                        text = '存在风险'
                        re_flag = True
                    else:
                        text = '无安全漏洞'
                else:
                    text = '未知'
            elif i==5 :
                if flag == 1 :text = u'无安全漏洞'
                if flag == 0 :text = '未知'  
            elif i==6:
                text = row[1]
            else:
                text = 0
                
            newItem = QtGui.QTableWidgetItem(unicode(str(text)))
            if re_flag: newItem.setBackgroundColor(QtGui.QColor(200,0,0))
            self.ui.result_tableWidget.setItem(self.count,i,newItem)
            
        self.count += 1
    def setTableRow(self):
        """
        爬虫结束后开始SQL注入检查
        """
        self.ui.result_tableWidget.scrollToBottom()
        self.logemit.emit(u"爬虫扫描结束")
        self.logemit.emit(u"去除重复相似页面共计:"+ str(UrlData.removSum))
        if not UrlData.end_flag:
            if UrlData.sqlflage:
                self.sql_map.start()
                self.sql_scan.start()
            elif UrlData.xssflage:
                self.xss_scan.start()
            else:
                self.switch()
        
    def insertResult(self):
        self.sql_result.start()

    def endSql(self):
        if UrlData.xssflage:
            self.xss_scan.start()
        else:
            self.showResult()

    def showResult(self):
        """
        检查解结束后,刷新表格
        """
        self.is_scanning = False
        self.switch()

        UrlData.end_flag = True
        count_tmp = self.count+1
        self.count = 0
        self.db_select.connectDb()
        cursor = self.db_select.selectAllInfo()
        lensum = self.ui.result_tableWidget.rowCount()
        for i in range(count_tmp):
            self.ui.result_tableWidget.removeRow(lensum)
            lensum -= 1
        for row in cursor:
            self.setTable(row,1)
        self.db_select.disConnect()
            
    #输出日至信息    
    def logSet(self,log_str):
        cursor = QtGui.QTextCursor
        cursor = self.ui.log_Edit.textCursor()
        cursor.movePosition(QtGui.QTextCursor.End)
        self.ui.log_Edit.setTextCursor(cursor)
        date_time = time.strftime("%H:%M:%S",time.localtime(time.time()))
        
        self.ui.log_Edit.insertPlainText(u'['+date_time+u'] '+log_str+'\n')
        cursor.movePosition(QtGui.QTextCursor.End)
        self.ui.log_Edit.setTextCursor(cursor)
        
    def __del__(self):
        self.spider.stop()
        UrlData.end_flag = True
        self.setting.hide()

    def switch(self):
        self.ui.startButton.setDisabled(self.is_scanning)
        self.ui.stopButton.setDisabled(not self.is_scanning)

    def initSetting(self):
        self.db_select.connectDb()
        cursor = self.db_select.conn.execute('select * from xss_pattern')
        self.setpyui.listWidget_rule.clear()

        for row in cursor:
            self.setpyui.listWidget_rule.addItem(str(row[0]) + " CODE: " + str(row[1])[:50] + ' RESULT: ' + str(row[2]))
        self.db_select.disConnect()

    def addXSSPattern(self):
        code = self.setpyui.lineEdit_code.text()
        result = self.setpyui.lineEdit_result.text()
        if code and result:
            try:
                self.db_select.connectDb()
                self.db_select.conn.execute('insert into xss_pattern (code, result) values(?, ?)', (str(code), str(result)))
                self.db_select.conn.commit()
                self.initSetting()
            except sqlite3.Error as e:
                self.logSet(e.message)
                self.db_select.conn.rollback()
                self.messageDialog(u'添加失败！')
            self.db_select.disConnect()

    def deleteXSSPattern(self):
        item = self.setpyui.listWidget_rule.currentItem()
        id = re.split(' CODE: | RESULT: ', unicode(item.text().toUtf8(), 'utf-8', 'ignore'))[0]
        try:
            self.db_select.connectDb()
            self.db_select.conn.execute('delete from xss_pattern where id=?',(int(id),))
            self.db_select.conn.commit()
            self.initSetting()
        except Exception as e:
            print e
            self.db_select.conn.rollback()
            self.messageDialog('删除失败！')
        self.db_select.disConnect()

    def modifyXSSPattern(self):
        item = self.setpyui.listWidget_rule.currentItem()
        id = re.split(' CODE: | RESULT: ', unicode(item.text().toUtf8(), 'utf-8', 'ignore'))[0]
        code = self.setpyui.lineEdit_code.text()
        result = self.setpyui.lineEdit_result.text()
        try:
            self.db_select.connectDb()
            self.db_select.conn.execute('update xss_pattern set code=?, result=? where id=?', (str(code), str(result), id))
            self.db_select.conn.commit()
            self.initSetting()
        except Exception as e:
            print e
            self.db_select.conn.rollback()
            self.messageDialog('修改失败！')
        self.db_select.disConnect()

    def _itemClicked(self, item):
        items = re.split(' CODE: | RESULT: ', unicode(item.text().toUtf8(), 'utf-8', 'ignore'))
        self.setpyui.lineEdit_code.setText(items[1])
        self.setpyui.lineEdit_result.setText(items[2])

    def messageDialog(self, msg, title=u'提示'):
        box = QtGui.QMessageBox(QtGui.QMessageBox.Warning, title, msg.decode('utf-8'), QtGui.QMessageBox.Yes)
        box.exec_()
        if box == QtGui.QMessageBox.Yes:
            box.hide()

#开始Sql注入检查 
class startSqlthread(threading.Thread):
    def __init__(self,url,id):
        threading.Thread.__init__(self)
        self.thread_stop=False
        self.url = url
        self.id = id
    def run(self):
        t = AutoStart('http://127.0.0.1:8775',self.url,self.id)
        get_id = t.run()
        if (get_id):
            UrlData.proce[self.id] = 1
        else:
            UrlData.proce[self.id] = 0
    def stop(self):
        self.thread_stop = True

class sqlmapThread(QtCore.QThread):
    logemit = QtCore.pyqtSignal(str)    
    def __int__(self):  
        super(tableThread,self).__init__()
    def run(self):
    
        self.startSqlmap()

    def startSqlmap(self):
        args_str = ["-s","-H","127.0.0.1","-p","8775"]
        apistart = apiStart()
        apistart.main(args_str)
    def osStart(self):
        # os.system("python2 ")
        pass
        
    
class sqlScanThread(QtCore.QThread):
    endemit = QtCore.pyqtSignal() 
    logemit = QtCore.pyqtSignal(str)
    labemit = QtCore.pyqtSignal(str)    
    def __int__(self):  
        super(sqlScanThread,self).__init__()
        
    def run(self):
        self.logemit.emit(u"开始Sql注入漏洞检查")
        self.labemit.emit(u"    Sql注入漏洞检查>>>>")
        time.sleep(5)       
        self.startSqlScan()    
    def startSqlScan(self):    
        #time.sleep(5)
        try:
            conn = sqlite3.connect('./data/urldata.db',check_same_thread = False)
        except sqlite3.Error,e:
            self.logemit.emit(u"连接数据库失败", "\n"+e.args[0])
        cursor = conn.execute("select url_full,id from get_url") 
        th_sum = threading.activeCount()
        for row in cursor:
            matchObj = re.search(r'(.html/|.html)$',row[0],re.I)
            while 1:
                if (threading.activeCount()-th_sum < 20):
                    break
            try:
                if not matchObj:
                    thread1=startSqlthread(row[0],row[1])
                    thread1.start()   
                else:
                    self.logemit.emit(u"忽略:"+row[0])
            except:
                self.logemit.emit(u"Error : unable to start thread")
        conn.close() 
        timeflage = 0
        while 1:
            self.labemit.emit(u"等待所有操作结束剩余线程数"+str(threading.activeCount()-1-th_sum))
            time.sleep(3)
            if (threading.activeCount() - th_sum <= 1 ): #and not UrlData.proce 
                break
            if (threading.activeCount() - th_sum <= 3 and timeflage <= 60):
                timeflage += 1
            if (timeflage >= 20): break
        #conn.close()
        self.logemit.emit(u"扫描结束")
        self.labemit.emit(u'SQL注入漏洞检查完毕')
        self.endemit.emit()
        # endemit
class insertSqlRuseltThread(QtCore.QThread):
    logemit = QtCore.pyqtSignal(str)
    endinsert = QtCore.pyqtSignal()
    def __int__(self):  
        super(insertSqlRuseltThread,self).__init__()
    def run(self):
        try:
            self.conn = sqlite3.connect('./data/urldata.db',check_same_thread = False)
        except sqlite3.Error,e:
            #print e.args[0]
            self.logemit.emit(u"数据库异常:"+e.args[0])
        self.deinnserst()
        self.endinsert.emit()
    def deinnserst(self):
        while  1:
            if UrlData.proce:
                for url_id,status in UrlData.proce.items(): 
                    self.conn.execute("insert into sql_result (url_id,status)values(?,?)",(int(url_id),status))
                    self.conn.commit()
                    #print url_id,status
                    del UrlData.proce[url_id]      
            else:
                break              
        self.conn.close()

    
#表格显示线程    
class tableThread(QtCore.QThread):  
    trigger = QtCore.pyqtSignal(tuple)  
    table = QtCore.pyqtSignal()
    def __int__(self):  
        super(tableThread,self).__init__()
        
  
    def run(self):

        self.unUrl = uniqueUrl()
        self.status = 0
        # self.db = dataBaseHelper()
        # self.db.open()  
        self.conn = sqlite3.connect('./data/urldata.db',check_same_thread = False)
        self.flage = 0
        while 1:
            count = 0
            while not UrlData.data_queue.empty() and count <= 300:
                node = UrlData.data_queue.get()
                url = node[1].url
                isrepeat = False
                try:
                    isrepeat = self.unUrl.add(url)
                except :
                    pass
                    #print "EOFError!!!!!!!!!!!!!!!!!"
                if  isrepeat and UrlData.removflage:
                    UrlData.removSum += 1                    
                else:
                    self.conn.execute("insert into get_url (url_full,method)values(?,?)",(node[1].url,node[1].method))
                    self.conn.commit()
                count += 1
            try:
                cursor = self.conn.execute("select id,url_full,add_date,status,method from get_url")
                if cursor:           
                    for row in cursor:
                        if self.status<row[0]:
                            self.trigger.emit(row)
                    self.status = row[0]
            except:
                time.sleep(3)
            if UrlData.spider_flage == 1:
                self.setTable()
                self.conn.close()
                return
            if UrlData.end_flag:
                self.setTable()
                self.conn.close()
                return
                
                   
    def setTable(self):
        if not UrlData.end_flag: self.table.emit()
        
#计划写一个数据库类
class dbClass(object):
    def __init__(self,dbpath='./data/urldata.db'):
        self.dbfile = dbpath
        
    def connectDb(self):
        try:
              self.conn = sqlite3.connect(self.dbfile,check_same_thread = False)
        except sqlite3.Error,e:
              print e.args[0]
    def selectAllInfo(self,ID=0):
        # if ID == 0:
        #     cursor = self.conn.execute("select get_url.id,get_url.url_full,get_url.add_date,sql_result.status ,get_url.method,get_url.status,sql_result.url_id from get_url,sql_result where sql_result.url_id=get_url.id")
        # else:
        #     cursor = self.conn.execute("select get_url.id,get_url.url_full,get_url.add_date,sql_result.status ,get_url.method,get_url.status,sql_result.url_id from get_url,sql_result where sql_result.url_id=? and get_url.id=?",(int(ID),int(ID)))
        if ID == 0:
            cursor = self.conn.execute("select get_url.id,get_url.url_full,get_url.add_date,sql_result.status, "
                                       "get_url.method,get_url.status,sql_result.url_id, xss_scan.result "
                                       "from get_url "
                                       "left join xss_scan on xss_scan.url_id=get_url.id "
                                       "left join sql_result on sql_result.url_id=get_url.id")
        else:
            cursor = self.conn.execute("select get_url.id,get_url.url_full,get_url.add_date,sql_result.status,"
                                       "get_url.method,get_url.status,sql_result.url_id, xss_scan.result "
                                       "from get_url "
                                       "left join xss_scan on xss_scan.url_id=?"
                                       "left join sql_result on sql_result.url_id=? "
                                       "where get_url.id=?",(int(ID),int(ID), int(ID)))
        return cursor
    def delectDb(self):
        self.conn.execute("DELETE FROM get_url WHERE 1=1")
        self.conn.commit()
        self.conn.execute("DELETE FROM sql_result WHERE 1=1 ")
        self.conn.commit()
        self.conn.execute("update sqlite_sequence set seq=0 where name='get_url'")
        self.conn.commit()
        self.conn.execute("DELETE FROM xss_scan")
        self.conn.execute("update sqlite_sequence set seq=0 where name='xss_scan'")
        self.conn.execute("update sqlite_sequence set seq=0 where name='sql_result'")
        self.conn.commit()
    def selectXSSpayload(self,ID=0):
        cursor = self.conn.execute("select patterns from xss_scan where url_id=?",(int(ID),))
        return cursor
    
    def disConnect(self):
        self.conn.close()
		
#Pyqt应用程序的主函数
if __name__ == "__main__":

    
    app = QtGui.QApplication(sys.argv)
    myapp = StartQT4()
    
    icon = QtGui.QIcon()
    icon.addPixmap(QtGui.QPixmap(u"./UIlib/enderman.ico"), QtGui.QIcon.Normal, QtGui.QIcon.Off)
    myapp.setWindowIcon(icon) 
    
    myapp.show()
    
    
    
    sys.exit(app.exec_())
    
    
    
        
    
 


